import torch
import torch.nn.functional as F
import torch.nn as nn

class LSTM(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim, layer_num):
        super(LSTM, self).__init__()
        self.input_dim = input_dim
        self.output_dim = output_dim
        self.hidden_dim = hidden_dim
        self.layer_num = layer_num

        self.lstm = nn.LSTM(self.input_dim, self.hidden_dim, self.layer_num, batch_first=True).cuda()

        self.fc = nn.Linear(self.hidden_dim, self.output_dim).cuda()
        self.fc1 = nn.Linear(self.output_dim, self.output_dim).cuda()
        self.dropout = nn.Dropout(p=0.9)

    def init_input(self, batch_size):
        h_0 = torch.randn(self.layer_num, batch_size, self.hidden_dim).cuda()
        c_0 = torch.randn(self.layer_num, batch_size, self.hidden_dim).cuda()
        return h_0, c_0

    def forward(self, x):
        batch_size = x.size(0)
        h_0, c_0 = self.init_input(batch_size)
        output, _ = self.lstm(x, (h_0, c_0))
        output = F.relu(output)
        output = self.dropout(output)
        output = self.fc(output[:, -1, :])
        output = nn.Tanh()(output)
        output = self.dropout(output)
        return output